# -*- coding: utf-8 -*-
import atexit
import sys
import ssl
import time

from pyVim.connect import Disconnect, SmartConnect
from pyVmomi import vim, vmodl

# vcenter信息
vcenter_ip = "172.16.100.57"
vcenter_port = "443"
vcenter_user = "Administrator@vsphere.local"
vcenter_pwd = "Shsnc123!@#"
vmware_name = "snc-test-imp"
flag = "delete"


# 虚拟机开关机需要的一个方法
def WaitForTasks(service_instance, tasks):
    """
    Given the service instance si and tasks, it returns after all the tasks are complete
    :param service_instance:
    :param tasks:
    :return:
    """
    pc = service_instance.content.propertyCollector
    task_list = [str(task) for task in tasks]
    # Create filter
    obj_specs = [vmodl.query.PropertyCollector.ObjectSpec(obj=task) for task in tasks]
    property_spec = vmodl.query.PropertyCollector.PropertySpec(type=vim.Task, pathSet=[], all=True)
    filter_spec = vmodl.query.PropertyCollector.FilterSpec()
    filter_spec.objectSet = obj_specs
    filter_spec.propSet = [property_spec]
    pcfilter = pc.CreateFilter(filter_spec, True)
    try:
        version, state = None, None
        # Loop looking for updates till the state moves to a completed state.
        while True:
            update = pc.WaitForUpdates(version)
            for filter_set in update.filterSet:
                for obj_set in filter_set.objectSet:
                    task = obj_set.obj
                    for change in obj_set.changeSet:
                        if change.name == 'info':
                            state = change.val.state
                        elif change.name == 'info.state':
                            state = change.val
                        else:
                            continue

            # Check if the state moved to a completed state.
            if state == vim.TaskInfo.State.success:
                return task.info.result
            elif state == vim.TaskInfo.State.error:
                raise task.info.error
    finally:
        if pcfilter:
            pcfilter.Destroy()


class vCenterManager():
    def __init__(self, vcenter_ip, vcenter_port, vcenter_user, vcenter_pwd):
        self.vcenter_ip = vcenter_ip
        self.vcenter_port = vcenter_port
        self.vcenter_user = vcenter_user
        self.vcenter_pwd = vcenter_pwd
        self.context = None
        self.si = None
        # 建立与vSphere的连接
        if hasattr(ssl, '_create_unverified_context'):
            self.context = ssl._create_unverified_context()
        self.si = SmartConnect(host=self.vcenter_ip, port=self.vcenter_port, user=self.vcenter_user,
                               pwd=self.vcenter_pwd,
                               sslContext=self.context)
        if self.si is None:
            print("请检查vc账号密码")
            sys.exit(1)

        atexit.register(Disconnect, self.si)

    def vm_power_oc(self, vm_name, flag):
        """
        虚拟机开关
        :param vm_name: 虚拟机名称
        :param flag: on/off: 开启、关闭
        :return:
        """
        content = self.si.RetrieveContent()
        # 查找虚拟机
        vm = None
        container = content.viewManager.CreateContainerView(content.rootFolder, [vim.VirtualMachine], True)
        vms = container.view
        for v in vms:
            if v.name == vm_name:
                vm = v
                break
        if vm:
            # 检查虚拟机当前状态
            if vm.runtime.powerState == "poweredOff" and flag == "on":
                # 开启虚拟机
                task = vm.PowerOnVM_Task()
                WaitForTasks(self.si, [task])
                print(f"虚拟机 '{vm_name}' 已开机")
            elif vm.runtime.powerState == "poweredOn" and flag == "off":
                # 关闭虚拟机
                task = vm.PowerOffVM_Task()
                WaitForTasks(self.si, [task])
                print(f"虚拟机 '{vm_name}' 已关闭")
            else:
                print(f"指令错误或找不到虚拟机 '{vm_name}'")

    # 虚拟机删除
    def delete_vm(self, vm_name):
        content = self.si.RetrieveContent()
        container = content.viewManager.CreateContainerView(content.rootFolder, [vim.VirtualMachine], True)
        vms = container.view
        for vm in vms:
            if vm.name == vm_name:
                try:
                    if vm.runtime.powerState == "poweredOn":
                        # 如果虚拟机正在运行，首先关闭它
                        task = vm.PowerOff()
                        task_info = task.info
                        task_info.state = vim.TaskInfo.State.success
                        task_info.result = None
                        # print(f"等待虚拟机 {vm_name} 关机...")
                        WaitForTasks(self.si, [task])
                    # 删除虚拟机
                    task = vm.Destroy()
                    task_info = task.info
                    task_info.state = vim.TaskInfo.State.success
                    task_info.result = None
                    start_time = time.time()
                    while task_info.state != vim.TaskInfo.State.success:
                        if time.time() - start_time > 300:  # 设置超时时间，单位秒
                            # raise Exception(f"删除虚拟机 {vm_name} 超时")
                            print({'code': 404, 'message': '虚拟机 %s 已超时5分钟' % vm_name})
                        time.sleep(5)  # 每5秒检查一次任务状态
                        task_info = task.info
                    print({'code': 200, 'message': '虚拟机 %s 已成功删除' % vm_name})
                except Exception as e:
                    print({'code': 404, 'message': ' %s' % e})


if __name__ == '__main__':
    vc = vCenterManager(vcenter_ip, vcenter_port, vcenter_user, vcenter_pwd)
    if flag == "on" or flag == "off":
        vc.vm_power_oc(vmware_name, flag)
    elif flag == "delete":
        vc.delete_vm(vmware_name)
    else:
        print("指令错误")